#!/usr/bin/env python

import glob
import optparse
import os
import sys

import tables

def main():
    op = optparse.OptionParser(usage="Usage: %prog [opts] nt_name output.h5 file1.h5 [file2.h5 [...]]")

    op.add_option("-z", "--compress", dest='compress', action='store_true', default=False,
                  help='Compress output file')
    op.add_option("-d", "--dir", dest="dir", action="store_true", default=False, 
                  help="Treat arguments as dirs and glob for HDF5 files")
    op.add_option("--file-list", dest="file_list", default=None, 
                  help="Use files listed in PATH for input", metavar="PATH")
    
    (opts, args) = op.parse_args()

    if len(args) < 2 and opts.file_list:
        sys.stderr.write("Please supply an ntuple name (PyTables dot notation ok) and an output file name.\n")
        sys.stderr.write( "See `h5append.py -h' for more.\n")
        sys.stderr.flush()
        return 1

    if len(args) < 3 and not opts.file_list:
        sys.stderr.write("Please supply an ntuple name, an output file, and at least one input file.\n")
        sys.stderr.write( "See `h5append.py -h' for more.\n")
        sys.stderr.flush()
        return 2

    ntname = args.pop(0)
    output  = args.pop(0)
    files_processed  = 0

    if opts.file_list:
        fp = open(opts.file_list, 'r')
        new_args = [line.strip() for line in fp]
        fp.close()

        new_args.sort()
        args = new_args
    elif opts.dir:
        newargs = []
        for path in args:
            newargs += [os.path.abspath(h5path) for h5path in glob.glob(os.path.join(path, '*.h5'))]

        args = newargs

    for path in args:
        if not tables.isHDF5File(path):
            sys.stderr.write("File `{0}' is not a valid HDF5 file. ".format(path))
            sys.stderr.write( "Omitting...\n")
            sys.stderr.flush()
            args.pop(0)
            continue
        elif not tables.isPyTablesFile(path):
            sys.stderr.write("File `{0}' is not a valid PyTables file. ".format(path))
            sys.stderr.write( "Omitting...\n")
            sys.stderr.flush()
            args.pop(0)
            continue

        first_file = args.pop(0)
        fpfirst = tables.openFile(first_file, 'r')
        
        try:
            ntfirst = eval("fpfirst.root.{0}".format(ntname))
            coldesc = ntfirst.coldescrs
            break
        except AttributeError:
            sys.stderr.write("File `{0}' does not contain an ntuple `{1}'. ".format(path, ntname))
            sys.stderr.write( "Omitting...\n")
            sys.stderr.flush()
            continue
            

    if output in args:
        args.remove(output)

    fpfirst.close()

    fpout = tables.openFile(output, 'w')

    if opts.compress:
        filt = tables.Filters(complevel=1, complib='lzo', shuffle=True)
    else:
        filt = None

    ntout = fpout.createTable(fpout.root, ntname, coldesc, 'created by h5append.py', filters=filt)
        
    args = [first_file] + args

    srcfiles = {}

    start_index = 0
    stop_index = 0
    for path in args:
        files_processed += 1

        sys.stderr.write("{0}/{1} files completed\r".format(files_processed, len(args))) 
        sys.stderr.flush()

        try:
            if not tables.isHDF5File(path):
                sys.stderr.write("\nFile `{0}' is not a valid HDF5 file. ".format(path))
                sys.stderr.write("Omitting...\n")
                sys.stderr.flush()
                continue
            if not tables.isPyTablesFile(path):
                sys.stderr.write("\nFile `{0}' is not a valid PyTables file. ".format(path))
                sys.stderr.write( "Omitting...\n")
                sys.stderr.flush()
                continue
            fpin = tables.openFile(path, mode='r')
            ntin = eval("fpin.root.{0}".format(ntname))

        except AttributeError:
            sys.stderr.write("\nFile `{0}' does not contain an ntuple `{1}'. ".format(path, ntname))
            sys.stderr.write( "Omitting...\n")
            sys.stderr.flush()
            fpin.close()
            continue
        except IOError, e:
            sys.stderr.write("\nFile `{0}' encountered an error:\n{1}\n".format(path, e))
            sys.stderr.write( "Omitting...\n")
            sys.stderr.flush()
            continue
        
        stop_index += len(ntin)

        ntout.append(ntin.read())
        ntout.flush()

        srcfiles[os.path.abspath(path)] = (start_index, stop_index)
        start_index = stop_index + 1

        fpin.close()

            
    print ''

    
    fpout.createArray(fpout.root, 'srcfiles', srcfiles.keys())
    fpout.createArray(fpout.root, 'srcfile_indicies', srcfiles.values())
    
    

    fpout.close()

if __name__ == "__main__":
    sys.exit(main())

